/*
 * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

package org.jetbrains.kotlin.test

import com.intellij.openapi.vfs.StandardFileSystems
import com.intellij.openapi.vfs.VirtualFileManager
import org.jetbrains.kotlin.cli.common.LegacyK2CliPipeline
import org.jetbrains.kotlin.cli.common.SessionWithSources
import org.jetbrains.kotlin.cli.common.messages.MessageCollector
import org.jetbrains.kotlin.cli.common.prepareJsSessions
import org.jetbrains.kotlin.cli.common.prepareJvmSessions
import org.jetbrains.kotlin.cli.common.prepareWasmSessions
import org.jetbrains.kotlin.cli.jvm.compiler.VfsBasedProjectEnvironment
import org.jetbrains.kotlin.cli.jvm.compiler.legacy.pipeline.MinimizedFrontendContext
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.fir.DependencyListForCliModule
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
import org.jetbrains.kotlin.fir.declarations.utils.nameOrSpecialName
import org.jetbrains.kotlin.fir.renderer.FirDeclarationRenderer
import org.jetbrains.kotlin.fir.renderer.FirDeclarationRendererWithFilteredAttributes
import org.jetbrains.kotlin.fir.renderer.FirRenderer
import org.jetbrains.kotlin.fir.resolve.providers.firProvider
import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
import org.jetbrains.kotlin.ir.backend.js.loadWebKlibsInTestPipeline
import org.jetbrains.kotlin.js.resolve.JsPlatformAnalyzerServices
import org.jetbrains.kotlin.library.loader.KlibPlatformChecker
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.platform.TargetPlatform
import org.jetbrains.kotlin.platform.isJs
import org.jetbrains.kotlin.platform.isWasm
import org.jetbrains.kotlin.platform.js.JsPlatforms
import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
import org.jetbrains.kotlin.platform.jvm.isJvm
import org.jetbrains.kotlin.platform.konan.isNative
import org.jetbrains.kotlin.platform.wasm.WasmPlatforms
import org.jetbrains.kotlin.platform.wasm.WasmTarget
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.resolve.DescriptorUtils
import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
import org.jetbrains.kotlin.test.directives.FirDiagnosticsDirectives
import org.jetbrains.kotlin.test.directives.FirDiagnosticsDirectives.PLATFORM_DEPENDANT_METADATA
import org.jetbrains.kotlin.test.directives.model.DirectivesContainer
import org.jetbrains.kotlin.test.directives.model.RegisteredDirectives
import org.jetbrains.kotlin.test.frontend.classic.ClassicFrontendOutputArtifact
import org.jetbrains.kotlin.test.frontend.fir.FirFrontendFacade
import org.jetbrains.kotlin.test.frontend.fir.FirOutputArtifact
import org.jetbrains.kotlin.test.frontend.fir.getAllJsDependenciesPaths
import org.jetbrains.kotlin.test.frontend.fir.getAllWasmDependenciesPaths
import org.jetbrains.kotlin.test.model.*
import org.jetbrains.kotlin.test.services.*
import org.jetbrains.kotlin.test.util.trimTrailingWhitespacesAndRemoveRedundantEmptyLinesAtTheEnd
import org.jetbrains.kotlin.test.utils.MultiModuleInfoDumper
import org.jetbrains.kotlin.test.utils.withExtension
import org.jetbrains.kotlin.utils.addToStdlib.shouldNotBeCalled
import org.jetbrains.kotlin.wasm.config.wasmTarget
import org.jetbrains.kotlin.wasm.resolve.WasmPlatformAnalyzerServices
import java.io.File

class JvmLoadedMetadataDumpHandler(testServices: TestServices) : AbstractLoadedMetadataDumpHandler<BinaryArtifacts.Jvm>(
    testServices,
    ArtifactKinds.Jvm
) {
    override val targetPlatform: TargetPlatform
        get() = JvmPlatforms.defaultJvmPlatform
    override val platformAnalyzerServices: PlatformDependentAnalyzerServices
        get() = JvmPlatformAnalyzerServices
    override val dependencyKind: DependencyKind
        get() = DependencyKind.Binary

    @OptIn(LegacyK2CliPipeline::class)
    override fun prepareSessions(
        module: TestModule,
        configuration: CompilerConfiguration,
        environment: VfsBasedProjectEnvironment,
        moduleName: Name,
        libraryList: DependencyListForCliModule,
    ): List<SessionWithSources<KtFile>> {
        return prepareJvmSessionsWithoutFiles(configuration, environment, moduleName, libraryList)
    }

    @LegacyK2CliPipeline
    private fun prepareJvmSessionsWithoutFiles(
        configuration: CompilerConfiguration,
        environment: VfsBasedProjectEnvironment,
        moduleName: Name,
        libraryList: DependencyListForCliModule
    ): List<SessionWithSources<KtFile>> {
        return MinimizedFrontendContext(environment, MessageCollector.NONE, emptyList(), configuration).prepareJvmSessions(
            files = emptyList(),
            moduleName,
            environment.getSearchScopeForProjectLibraries(),
            libraryList,
            isCommonSource = { false },
            isScript = { false },
            fileBelongsToModule = { _, _ -> false },
            createProviderAndScopeForIncrementalCompilation = { null }
        )
    }

}

class KlibJsLoadedMetadataDumpHandler(testServices: TestServices) : AbstractLoadedMetadataDumpHandler<BinaryArtifacts.KLib>(
    testServices,
    ArtifactKinds.KLib
) {
    override val targetPlatform: TargetPlatform
        get() = JsPlatforms.defaultJsPlatform
    override val platformAnalyzerServices: PlatformDependentAnalyzerServices
        get() = JsPlatformAnalyzerServices
    override val dependencyKind: DependencyKind
        get() = DependencyKind.Binary

    override fun prepareSessions(
        module: TestModule,
        configuration: CompilerConfiguration,
        environment: VfsBasedProjectEnvironment,
        moduleName: Name,
        libraryList: DependencyListForCliModule,
    ): List<SessionWithSources<KtFile>> {
        val klibs = loadWebKlibsInTestPipeline(
            configuration = configuration,
            libraryPaths = getAllJsDependenciesPaths(module, testServices),
            platformChecker = KlibPlatformChecker.JS,
        )

        return prepareJsSessions(
            files = emptyList(),
            configuration,
            moduleName,
            klibs.all,
            libraryList,
            extensionRegistrars = emptyList(),
            isCommonSource = { false },
            fileBelongsToModule = { _, _ -> false },
            icData = null
        )
    }
}

class KlibWasmJsLoadedMetadataDumpHandler(testServices: TestServices) : AbstractLoadedMetadataDumpHandler<BinaryArtifacts.KLib>(
    testServices,
    ArtifactKinds.KLib
) {
    override val targetPlatform: TargetPlatform
        get() = WasmPlatforms.wasmJs
    override val platformAnalyzerServices: PlatformDependentAnalyzerServices
        get() = WasmPlatformAnalyzerServices
    override val dependencyKind: DependencyKind
        get() = DependencyKind.Binary

    override fun prepareSessions(
        module: TestModule,
        configuration: CompilerConfiguration,
        environment: VfsBasedProjectEnvironment,
        moduleName: Name,
        libraryList: DependencyListForCliModule,
    ): List<SessionWithSources<KtFile>> {
        val klibs = loadWebKlibsInTestPipeline(
            configuration = configuration,
            libraryPaths = getAllWasmDependenciesPaths(module, testServices, configuration.wasmTarget),
            platformChecker = KlibPlatformChecker.Wasm(WasmTarget.JS.alias),
        )

        return prepareWasmSessions(
            files = emptyList(),
            configuration,
            moduleName,
            klibs.all,
            libraryList,
            extensionRegistrars = emptyList(),
            isCommonSource = { false },
            fileBelongsToModule = { _, _ -> false },
            icData = null
        )
    }
}

abstract class AbstractLoadedMetadataDumpHandler<A : ResultingArtifact.Binary<A>>(
    testServices: TestServices,
    override val artifactKind: ArtifactKind<A>
) : BinaryArtifactHandler<A>(
    testServices,
    artifactKind,
    failureDisablesNextSteps = false,
    doNotRunIfThereWerePreviousFailures = false
) {
    private val dumper: MultiModuleInfoDumper = MultiModuleInfoDumper()

    override val directiveContainers: List<DirectivesContainer>
        get() = listOf(FirDiagnosticsDirectives)

    override fun processModule(module: TestModule, info: A) {
        val languageSettingsBuilder = testServices.defaultsProvider.newLanguageSettingsBuilder()
        languageSettingsBuilder.configureUsingDirectives(
            testServices.defaultDirectives,
            testServices.environmentConfigurators,
            testServices.defaultsProvider.frontendKind == FrontendKinds.FIR
        )

        val emptyModule = TestModule(
            name = "dump-${module.name}", files = emptyList(),
            allDependencies = listOf(DependencyDescription(module, dependencyKind, DependencyRelation.RegularDependency)),
            RegisteredDirectives.Empty, languageSettingsBuilder.build()
        )
        val configuration = testServices.compilerConfigurationProvider.getCompilerConfiguration(emptyModule)
        val environment = VfsBasedProjectEnvironment(
            testServices.compilerConfigurationProvider.getProject(emptyModule),
            VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL),
            testServices.compilerConfigurationProvider.getPackagePartProviderFactory(emptyModule)
        )
        val moduleName = Name.identifier(emptyModule.name)
        val libraryList = FirFrontendFacade.initializeLibraryList(
            emptyModule, moduleName, targetPlatform, configuration, testServices
        )

        val session = prepareSessions(
            emptyModule,
            configuration,
            environment,
            moduleName,
            libraryList
        ).single().session

        val packageFqName = FqName("test")
        val printAttributes = FirDiagnosticsDirectives.RENDER_FIR_DECLARATION_ATTRIBUTES in module.directives
        dumper.builderForModule(module)
            .append(collectPackageContent(session, packageFqName, extractNames(module, packageFqName), printAttributes))
    }

    protected abstract val targetPlatform: TargetPlatform
    protected abstract val platformAnalyzerServices: PlatformDependentAnalyzerServices
    protected abstract val dependencyKind: DependencyKind

    protected abstract fun prepareSessions(
        module: TestModule,
        configuration: CompilerConfiguration,
        environment: VfsBasedProjectEnvironment,
        moduleName: Name,
        libraryList: DependencyListForCliModule,
    ): List<SessionWithSources<KtFile>>

    override fun processAfterAllModules(someAssertionWasFailed: Boolean) {
        if (dumper.isEmpty()) return
        val testDataFile = testServices.moduleStructure.originalTestDataFiles.first()

        val frontendKind = testServices.defaultsProvider.frontendKind

        val commonExtension = ".fir.txt"
        val (specificExtension, otherSpecificExtension) = when (frontendKind) {
            FrontendKinds.ClassicFrontend -> ".fir.k1.txt" to ".fir.k2.txt"
            FrontendKinds.FIR -> ".fir.k2.txt" to ".fir.k1.txt"
            else -> shouldNotBeCalled()
        }

        val targetPlatform = testServices.moduleStructure.modules.last().targetPlatform(testServices)
        if (PLATFORM_DEPENDANT_METADATA in testServices.moduleStructure.allDirectives) {
            val platformExtension = specificExtension.replace(".txt", "${targetPlatform.suffix}.txt")
            val otherPlatformExtension = specificExtension.replace(".txt", "${targetPlatform.oppositeSuffix}.txt")

            val expectedFile = testDataFile.withExtension(platformExtension)
            val actualText = dumper.generateResultingDump()
            assertions.assertEqualsToFile(expectedFile, actualText, message = { "Content is not equal" })

            val checks = listOf(commonExtension, specificExtension, otherSpecificExtension).map { extension ->
                {
                    val baseFile = testDataFile.withExtension(extension)
                    assertions.assertFalse(baseFile.exists()) {
                        "Base file $baseFile exists in presence of $PLATFORM_DEPENDANT_METADATA directive. Please remove file or directive"
                    }
                }
            }
            assertions.assertAll(checks)
            val secondFile = testDataFile.withExtension(otherPlatformExtension)
            val common = testDataFile.withExtension(specificExtension)
            checkDumpsIdentity(
                testDataFile, expectedFile, secondFile, common,
                postProcessTestData = { it.replace("// $PLATFORM_DEPENDANT_METADATA\n", "") }
            )
        } else {
            val commonFirDump = testDataFile.withExtension(commonExtension)
            val specificFirDump = testDataFile.withExtension(specificExtension)

            val expectedFile = when {
                commonFirDump.exists() -> commonFirDump
                else -> specificFirDump
            }

            val actualText = dumper.generateResultingDump()
            assertions.assertEqualsToFile(expectedFile, actualText, message = { "Content is not equal" })


            if (commonFirDump.exists() && specificFirDump.exists()) {
                assertions.fail {
                    """
                    Common dump ${commonFirDump.name} and specific ${specificFirDump.name} exist at the same time
                    Please remove ${specificFirDump.name}
                """.trimIndent()
                }
            }
            if (!commonFirDump.exists()) {
                val otherFirDump = testDataFile.withExtension(otherSpecificExtension)
                checkDumpsIdentity(testDataFile, specificFirDump, otherFirDump, commonFirDump)
            }
        }
    }

    private val TargetPlatform.suffix: String
        get() = when {
            isJvm() -> ".jvm"
            isJs() -> ".klib"
            isWasm() -> ".klib"
            isNative() -> ".klib"
            else -> error("Unsupported platform: $this")
        }

    private val TargetPlatform.oppositeSuffix: String
        get() = when {
            isJvm() -> ".klib"
            isJs() -> ".jvm"
            isWasm() -> ".jvm"
            isNative() -> ".jvm"
            else -> error("Unsupported platform: $this")
        }

    private fun checkDumpsIdentity(
        testDataFile: File,
        file1: File,
        file2: File,
        commonFile: File,
        postProcessTestData: ((String) -> String)? = null
    ) {
        if (!file1.exists() || !file2.exists()) return
        val dump1 = file1.readText().trimTrailingWhitespacesAndRemoveRedundantEmptyLinesAtTheEnd()
        val dump2 = file2.readText().trimTrailingWhitespacesAndRemoveRedundantEmptyLinesAtTheEnd()
        if (dump1 == dump2) {
            commonFile.writeText(dump1)
            file1.delete()
            file2.delete()
            if (postProcessTestData != null) {
                testDataFile.writeText(postProcessTestData(testDataFile.readText()))
            }
            assertions.fail {
                """
                    Files ${file1.name} and ${file2.name} are identical
                    Generating ${commonFile.name} and deleting ${file1.name} and ${file2.name}
                """.trimIndent()
            }
        }
    }

    private fun collectPackageContent(session: FirSession, packageFqName: FqName, declarationNames: Collection<Name>, printAttributes: Boolean): String {
        val provider = session.symbolProvider

        val builder = StringBuilder()
        val firRenderer = FirRenderer(
            builder,
            declarationRenderer = if (printAttributes) {
                FirDeclarationRendererWithFilteredAttributes()
            } else {
                FirDeclarationRenderer()
            }
        )

        for (name in declarationNames) {
            for (symbol in provider.getTopLevelCallableSymbols(packageFqName, name)) {
                firRenderer.renderElementAsString(symbol.fir)
                builder.appendLine()
            }
        }

        for (name in declarationNames) {
            val classLikeSymbol = provider.getClassLikeSymbolByClassId(ClassId.topLevel(packageFqName.child(name))) ?: continue
            firRenderer.renderElementAsString(classLikeSymbol.fir)
            builder.appendLine()
        }

        return builder.toString().trimEnd()
    }

    private fun extractNames(module: TestModule, packageFqName: FqName): Collection<Name> {
        testServices.artifactsProvider.getArtifactSafe(module, FrontendKinds.ClassicFrontend)
            ?.let { return extractNames(it, packageFqName) }
        testServices.artifactsProvider.getArtifactSafe(module, FrontendKinds.FIR)
            ?.let { return extractNames(it, packageFqName) }
        error("Frontend artifact for module $module not found")
    }

    private fun extractNames(artifact: ClassicFrontendOutputArtifact, packageFqName: FqName): Collection<Name> {
        return DescriptorUtils.getAllDescriptors(artifact.analysisResult.moduleDescriptor.getPackage(packageFqName).memberScope)
            .mapTo(sortedSetOf()) { it.name }
    }

    private fun extractNames(artifact: FirOutputArtifact, packageFqName: FqName): Collection<Name> {
        return sortedSetOf<Name>().apply {
            for (part in artifact.partsForDependsOnModules) {
                val files = part.session.firProvider.getFirFilesByPackage(packageFqName)
                files.flatMapTo(this) { file ->
                    file.declarations.mapNotNull { (it as? FirMemberDeclaration)?.nameOrSpecialName }
                }
            }
        }
    }
}
